#include "GCM.h"
#include "array/Arrays.h"
#include "lang/System.h"
#include "crypto/CryptoUtil.h"

using namespace yzrilyzr_array;
using namespace yzrilyzr_lang;
namespace yzrilyzr_crypto{
	GCM::GCM(const ByteArray & iv, const ByteArray & aad, u_index tagLength){
		if(iv == nullptr || iv.length == 0) throw IllegalArgumentException("IV cannot be null");
		this->iv=iv.clone();
		this->aad=aad != nullptr?aad.clone():ByteArray(0);
		this->tagLength=tagLength;
	}
	void GCM::init(FAQCipherAlgorithm & algorithm){
		// 1. 计算哈希子密钥H = E(K, 0^128)
		// 哈希子密钥
		static yzrilyzr_array::ByteArray ZERO_BLOCK(16);
		ByteArray h=algorithm.encrypt(ZERO_BLOCK);
		ghash.reset();
		ghash.setKey(FAQSecretKey(h, ""));
		// 2. 生成J0
		if(iv.length == 12){
			J0=ByteArray(16);
			System::arraycopy(iv, 0, J0, 0, 12);
			J0[15]=0x01; // J0 = IV || 0x00000001
		} else{
			// J0 = GHASH(H, {}, IV)
			ByteArray paddedIV=padBlock(iv, 16);
			ghash.updateAAD(ByteArray(0)); // 空AAD
			ghash.update(paddedIV);
			J0=ghash.doFinal();
		}
		Y_Inc_Cache=J0.clone();
		initialized=true;
	}

	ByteArray GCM::padBlock(const ByteArray & input, u_index blockSize){
		if(input.length % blockSize == 0) return input.clone();
		u_index paddedLen=((input.length + blockSize - 1) / blockSize) * blockSize;
		ByteArray padded(paddedLen);
		System::arraycopy(input, 0, padded, 0, input.length);
		return padded;
	}

	ByteArray GCM::encrypt(const ByteArray & block, FAQCipherAlgorithm & algorithm){
		checkInitialized();
		buffer.write(block, 0, block.length);
		return ByteArray(0);
	}
	ByteArray GCM::encryptFinal(const ByteArray & block, FAQCipherAlgorithm & algorithm){
		checkInitialized();
		buffer.write(block, 0, block.length);
		ByteArray P=buffer.toByteArray();
		buffer.reset();
		copyCounterToCache();
		increment32ToThis(Y_Inc_Cache);
		ByteArray C=GCTR(algorithm, Y_Inc_Cache, P);
		ghash.updateAAD(aad);
		ghash.update(C);
		ByteArray S=ghash.doFinal();
		ByteArray encryptedS=GCTR(algorithm, J0, S);
		ByteArray T=Arrays::copyOf(encryptedS, tagLength);
		return yzrilyzr_util::Util::concat(C, T);
	}
	ByteArray GCM::decrypt(const ByteArray & block, FAQCipherAlgorithm & algorithm){
		checkInitialized();
		buffer.write(block, 0, block.length);
		return ByteArray(0);
	}
	ByteArray GCM::decryptFinal(const ByteArray & block, FAQCipherAlgorithm & algorithm){
		checkInitialized();
		buffer.write(block, 0, block.length);
		ByteArray ciphertextWithTag=buffer.toByteArray();
		buffer.reset();
		if(ciphertextWithTag.length < tagLength) throw IllegalArgumentException(u"密文长度不足");
		u_index ciphertextLen=ciphertextWithTag.length - tagLength;
		ByteArray C=Arrays::copyOfRange(ciphertextWithTag, 0, ciphertextLen);
		ByteArray receivedTag=Arrays::copyOfRange(ciphertextWithTag, ciphertextLen, ciphertextWithTag.length);
		copyCounterToCache();
		increment32ToThis(Y_Inc_Cache);
		ByteArray P=GCTR(algorithm, Y_Inc_Cache, C);
		ghash.updateAAD(aad);
		ghash.update(C);
		ByteArray S=ghash.doFinal();
		ByteArray encryptedS=GCTR(algorithm, J0, S);
		ByteArray T=Arrays::copyOf(encryptedS, tagLength);
		if(!CryptoUtil::constantTimeEquals(T, receivedTag))
			throw SecurityException(u"认证标签验证失败");
		// 解密密文
		return P;
	}
	void GCM::copyCounterToCache(){
		System::arraycopy(J0, 0, Y_Inc_Cache, 0, J0.length);
	}
	ByteArray GCM::increment32ToThis(ByteArray value){
	   // 从最后字节开始，只处理最后4个字节(32位)
		u_index n=value.length - 1;
		while((n >= value.length - 4) && ((++value[n]) == 0)){
			n--;
		}
		return value;
	}
	ByteArray GCM::GCTR(FAQCipherAlgorithm & algorithm, ByteArray & ctr, ByteArray & input){
		ByteArray output(input.length);
		for(u_index i=0;i < input.length;i+=16){
			ByteArray encrypted=algorithm.encrypt(ctr);
			u_index blockLen=Math::min((u_index)16, input.length - i);
			for(u_index j=0;j < blockLen;j++){
				output[i + j]=(int8_t)(input[i + j] ^ encrypted[j]);
			}
			increment32ToThis(ctr);
		}
		return output;
	}

	void GCM::checkInitialized(){
		if(!initialized) throw IllegalStateException(u"GCM未初始化");
	}
	GHash & GCM::getGHash(){
		return ghash;
	}
}